# Projection-Iterative-Methods-based Optimizer
# https://www.sciencedirect.com/science/article/pii/S0950705125010238

import numpy as np

from scipy.special import gamma

from util.util_log import setup_logging
from util.util_csv import save_csv, save_to_csv
from util.util_image import save_img, save_img_xyz, save_line_chart
from util.util_ris_pattern_2 import point_2_phi_pattern, phase_2_bit, phase_2_pattern, phase_2_pattern_xyz
from beam.util_psll import get_psll_by_phase, get_psll_by_phaseBit


class RISProjIterMethOptimizer:
    __bit_num = 0  # 比特数
    __beam_num = 0  # 波束数
    __ub = 0
    __lb = 1

    def __init__(self, bit_num, beam_num, population_size=50, num_generations=100):
        # 初始化阵列相关
        self.__bit_num = bit_num
        self.__beam_num = beam_num
        # 初始化算法相关
        self.population_size = population_size
        self.num_generations = num_generations
        self.population = None
        self.best_individual = None
        self.best_fitness = None
        self.best_fitness_history = []  # 保存每一代的最佳适应度值
        self.best_individual_history = []  # 保存每一代的最佳适应度值的个体

    def initialize_population(self, phase_mix_init):
        """初始化种群"""
        # 扁平化为一维数组
        flat_phase_mix_init = phase_mix_init.flatten()
        self.population = [flat_phase_mix_init] * self.population_size
        return len(flat_phase_mix_init)

    def fitness(self, phase):
        return get_psll_by_phase(phase, self.__bit_num, self.__beam_num)

    def get_jacobian(self, x, epsilon=1e-6):
        """
        使用有限差分法计算广义雅可比矩阵
        参数:
            x: 当前点（向量）
            epsilon: 小扰动值
        返回:
            J: 雅可比矩阵
        """
        n = len(x)
        fx = self.fitness(x)
        # 如果输出是标量，将其转换为数组
        if np.isscalar(fx) or (hasattr(fx, '__len__') and len(np.array(fx).shape) == 0):
            fx = np.array([fx])
            m = 1
        else:
            fx = np.array(fx)
            m = len(fx)
        #
        J = np.zeros((m, n))
        #
        # 计算每个方向的有限差分
        for i in range(n):
            x_perturbed = x.copy()
            x_perturbed[i] = x_perturbed[i] + epsilon
            f_perturbed = self.fitness(x_perturbed)
            # 如果输出是标量，转换为数组
            if np.isscalar(f_perturbed) or (hasattr(f_perturbed, '__len__') and len(np.array(f_perturbed).shape) == 0):
                f_perturbed = np.array([f_perturbed])
            else:
                f_perturbed = np.array(f_perturbed)
            J[:, i] = (f_perturbed - fx) / epsilon
        #
        return J


    def levy_flight(self, n, m, beta=1.5):
        """
        莱维飞行分布生成函数
        参数:
            n: 行数
            m: 列数
            beta: 莱维分布参数
        返回:
            z: 莱维飞行步长矩阵
        """
        # 计算莱维分布参数
        num = gamma(1 + beta) * np.sin(np.pi * beta / 2)
        den = gamma((1 + beta) / 2) * beta * (2 ** ((beta - 1) / 2))
        sigma_u = (num / den) ** (1 / beta)
        # 生成正态分布随机数
        u = np.random.normal(0, sigma_u, (n, m))
        v = np.random.normal(0, 1, (n, m))
        # 计算莱维飞行步长
        z = u / (np.abs(v) ** (1 / beta))
        #
        return z

    def run(self, phase_mix_init, logger):
        logger.info("population_size=%d, num_generations=%d" % (self.population_size, self.num_generations))
        epsilon = 1e-6
        alphabet = np.arange(1, self.population_size + 1)

        # """初始化返回值"""
        self.best_fitness = self.fitness(phase_mix_init)
        self.best_individual = phase_mix_init
        self.best_fitness_history = []
        self.best_individual_history = []

        # """运行算法 -- 初始化阶段"""
        dim = self.initialize_population(phase_mix_init)

        # 适应度列表
        fitness = np.array([self.fitness(ind) for ind in self.population])
        sorted_indices = np.argsort(fitness)
        self.best_fitness = fitness[sorted_indices[0]]
        self.best_individual = self.population[sorted_indices[0], :].copy()

        logger.info(f"迭代 0: 最优适应度 = {self.best_fitness:.6f}")

        # """运行算法 -- 搜索阶段"""
        for it in range(self.num_generations):
            # 计算delta参数，用于控制步长
            delta = np.sin(np.pi / 2 * (1 - (2 * it / self.num_generations)) ** 5)
            a = (1 - it / self.num_generations) * np.random.rand(dim)

            for i in range(self.population_size):
                rr1 = np.random.rand()
                rr2 = np.random.rand()

                if rr1 >= rr2:
                    # 基于适应度的概率选择
                    pr = fitness ** 2 / (np.linalg.norm(fitness) ** 2 + np.finfo(float).eps)
                    try:
                        # 使用概率选择个体
                        p_indices = np.random.choice(alphabet, size=2, p=pr / np.sum(pr))
                        p1, p2 = int(p_indices[0]) - 1, int(p_indices[1]) - 1  # 转换为0基索引
                    except:
                        # 如果概率选择失败，随机选择
                        p = np.random.randint(0, self.population_size, 2)
                        p1, p2 = p[0], p[1]
                else:
                    # 选择前两个最优个体
                    p = sorted_indices[:2]
                    p1, p2 = p[0], p[1]

                r = np.random.rand()

                # 计算广义雅可比矩阵
                jacobian = self.get_jacobian(self.population[p1, :], epsilon)

                if 3 * np.random.rand() >= 2 * np.random.rand():
                    # 公式(14): 使用梯度下降
                    grad = (r * (self.population[p1, :] - self.best_individual) +
                            (1 - r) * (self.population[p2, :] - self.best_individual)) / 2
                    pos_n1 = self.population[i, :] - delta * grad
                else:
                    # 公式(15): 使用雅可比矩阵
                    grad = (r * (self.population[p2, :] - self.best_individual) +
                            (1 - r) * (self.population[p1, :] - self.best_individual)) / 2
                    # 如果雅可比矩阵是1维输出，则直接使用梯度
                    if jacobian.shape[0] == 1:
                        pos_n1 = self.population[i, :] + delta * jacobian.flatten() * grad
                    else:
                        pos_n1 = self.population[i, :] + delta * jacobian.T @ grad

                # 边界处理
                newpos1 = np.clip(pos_n1, self.__lb, self.__ub)
                fitt = self.fitness(newpos1)

                if fitt < fitness[i]:
                    fitness[i] = fitt
                    self.population[i, :] = newpos1

                # 额外的位置更新步骤
                if np.random.rand() > np.random.rand():
                    p = np.random.permutation(self.population_size)[:2]  # 公式(13)
                    p1, p2 = p[0], p[1]

                    if 3 * np.random.rand() >= 2 * np.random.rand():
                        # 公式(14)
                        grad = (r * (self.population[p1, :] - self.best_individual) +
                                (1 - r) * (self.population[p2, :] - self.best_individual)) / 2
                        pos_n2 = self.population[i, :] - delta * grad
                    else:
                        # 公式(15)
                        grad = (r * (self.population[p2, :] - self.best_individual) +
                                (1 - r) * (self.population[p1, :] - self.best_individual)) / 2
                        # 如果雅可比矩阵是1维输出，则直接使用梯度
                        if jacobian.shape[0] == 1:
                            pos_n2 = self.population[i, :] + delta * jacobian.flatten() * grad
                        else:
                            pos_n2 = self.population[i, :] + delta * jacobian.T @ grad

                    # 边界处理
                    newpos2 = np.clip(pos_n2, self.__lb, self.__ub)
                    fitt = self.fitness(newpos2)

                    if fitt < fitness[i]:
                        fitness[i] = fitt
                        self.population[i, :] = newpos2

                # 维度级别的更新
                for j in range(dim):
                    r1 = 1 + np.random.rand()  # 公式(16)
                    r2 = 1 + np.random.rand()

                    # 公式(17)
                    pho_1 = r1 * self.population[i, :] + (1 - r1) * self.best_individual \
                            + r2 * (self.population[i, :] - self.best_individual)
                    # 公式(18)
                    pho_2 = self.population[i, :] + a * (newpos1 - self.best_individual)

                    pos_n3 = self.population[i, :].copy()

                    if np.random.rand() / (j + 1) > np.random.rand():
                        pos_n3[j] = pho_1[j]
                    else:
                        pos_n3[j] = pho_2[j]

                    # 边界处理
                    newpos3 = np.clip(pos_n3, self.__lb, self.__ub)
                    fitt = self.fitness(newpos3)

                    if fitt < fitness[i]:
                        fitness[i] = fitt
                        self.population[i, :] = newpos3

            # 莱维飞行阶段
            if np.random.rand() <= 0.5 * (np.tanh(9 * it / self.num_generations - 5) + 1):
                pos_n4 = np.zeros(dim)
                d = np.random.rand() * (1 - it / self.num_generations) ** 2
                step_length = self.levy_flight(self.population_size, dim, 1.5)
                elite = np.tile(self.best_individual, (self.population_size, 1))
                r = np.random.rand()

                for i in range(self.population_size):
                    for j in range(dim):
                        pos_n4[j] = (r * elite[i, j] +
                                     (1 - r) * step_length[i, j] * d *
                                     (elite[i, j] - self.population[i, j] * (2 * it / self.num_generations)))  # 公式(21)

                    # 边界处理
                    newpos4 = np.clip(pos_n4, self.__lb, self.__ub)
                    fitt = self.fitness(newpos4)

                    if fitt < fitness[i]:
                        fitness[i] = fitt
                        self.population[i, :] = newpos4

            # 记录收敛曲线
            sorted_indices = np.argsort(fitness)
            if fitness[sorted_indices[0]] < self.best_fitness:
                self.best_fitness = fitness[sorted_indices[0]]
                self.best_individual = self.population[sorted_indices[0], :].copy()

            # 记录历史数据
            self.best_fitness_history.append(self.best_fitness)
            self.best_individual_history.append(self.best_individual)

            logger.info(f"迭代 {it + 1}: 最优适应度 = {self.best_fitness:.6f}")

        return self.best_individual, self.best_fitness, self.best_fitness_history, self.best_individual_history



# ============================================= 主函数 ====================================
# 用 phase 计算 phaseBit & patternBit
def get_bit_pattern(phase, bit_num):
    # 相位转换 X bit
    phaseBit, phaseBitDeg = phase_2_bit(phase, bit_num)
    # 计算phase_mix的方向图
    phaseBit = np.deg2rad(phaseBitDeg)
    patternBit = phase_2_pattern(phaseBit)
    return phaseBit, patternBit


# 核心方法 -- 单波束 -- PIMO
def pimo_beam1(phase, bit_num, pimo):
    # 2.PIMO寻找最优的phaseBit
    best_individual, best_fitness, best_fitness_history, best_individual_history = pimo.run(phase, logger)
    logger.info("PSO.best_fitness=%f" % (best_fitness))
    # 3.计算优化后的pattern
    phase_pimo = best_individual
    phaseBit_pimo, patternBit_pimo = get_bit_pattern(phase_pimo, bit_num)
    #
    return phase_pimo, phaseBit_pimo, patternBit_pimo, best_fitness, best_fitness_history, best_individual_history


# 单波束 -- 副瓣抑制 -- PIMO
def main_beam1_psll(theta_start, theta_step, theta_end,
                    phi_start, phi_step, phi_end,
                    bit_num=1):
    logger.info(f"main_beam1_psll: theta_start={theta_start}, theta_step={theta_step}, theta_end={theta_end}")
    logger.info(f"main_beam1_psll: phi_start={phi_start}, phi_step={phi_step}, phi_end={phi_end}")
    max_iter = 50
    beam_num = 1
    #
    pimo = RISProjIterMethOptimizer(bit_num=bit_num, beam_num=beam_num, population_size=50, num_generations=max_iter)
    #
    for theta in range(theta_start, theta_end, theta_step):
        for phi in range(phi_start, phi_end, phi_step):
            logger.info(f"theta:{theta}, phi:{phi}")
            # 1.背投影法(BP)计算码阵
            phase, phaseBit, patternBit = point_2_phi_pattern(theta, phi, bit_num)
            # 2.PIMO算法(PIMO)优化
            phase_pimo, phaseBit_pimo, patternBit_pimo, \
            best_fitness_pimo, best_fitness_history_pimo, best_individual_history_pimo \
                = pimo_beam1(phase, bit_num, pimo)
            # 计算pattern
            patternBit_xyz, x, y, z = phase_2_pattern_xyz(phaseBit)
            patternBit_xyz_pimo, x_pimo, y_pimo, z_pimo = phase_2_pattern_xyz(phaseBit_pimo)
            # 保存结果
            logger.info("save beam 1 result...")
            # 保存指标
            psll = get_psll_by_phaseBit(phaseBit, bit_num, 1)
            psll_pimo = get_psll_by_phaseBit(phaseBit_pimo, bit_num, 1)
            print(f"psll: {psll}, psll_pimo: {psll_pimo}")


# 使用示例
if __name__ == "__main__":
    # 初始化日志
    logger = setup_logging()
    # logger = setup_logging(base_path + "/log.txt")
    # 日志记录
    theta_start, theta_step, theta_end, phi_start, phi_step, phi_end = 30, 10, 45, 0, 10, 5
    logger.info(f"theta_start={theta_start}, theta_step={theta_step}, theta_end={theta_end}")
    logger.info(f"phi_start={phi_start}, phi_step={phi_step}, phi_end={phi_end}")

    main_beam1_psll(theta_start, theta_step, theta_end, phi_start, phi_step, phi_end)



